# src/gui/guipc.py

import os
import sys
import traceback
import tkinter as tk
from tkinter import filedialog, messagebox

import pandas as pd

from src.calphad import create_calculator
from src.calphad import CalphadPy


class CalphadGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("PyCalphad Equilibrium Linear Calculator")
        self.selected_elements = {}  # dict{elem: fracts(float)}
        self.element_buttons = {}   # dict{elem: ButtonWidget}
        self.phase_vars = {}        # dict{phase: BooleanVar}
        self.available_species = []  # list[elem] in tdb

        # periodic data added VA, e⁻
        self.periodic_table_data = {
            "elements": {
                -1: ("e⁻", "Electron", "special"),  # electron set Atomic number as -1
                0: ("VA", "Vacancy", "special"),    # vacancy set atomic number as 0
                1: ("H", "Hydrogen", "nonmetal"),
                2: ("He", "Helium", "noble gas"),

                3: ("Li", "Lithium", "alkali metal"),
                4: ("Be", "Beryllium", "alkaline earth metal"),
                5: ("B", "Boron", "metalloid"),
                6: ("C", "Carbon", "nonmetal"),
                7: ("N", "Nitrogen", "nonmetal"),
                8: ("O", "Oxygen", "nonmetal"),
                9: ("F", "Fluorine", "nonmetal"),
                10: ("Ne", "Neon", "noble gas"),

                11: ("Na", "Sodium", "alkali metal"),
                12: ("Mg", "Magnesium", "alkaline earth metal"),
                13: ("Al", "Aluminum", "post-transition metal"),
                14: ("Si", "Silicon", "metalloid"),
                15: ("P", "Phosphorus", "nonmetal"),
                16: ("S", "Sulfur", "nonmetal"),
                17: ("Cl", "Chlorine", "nonmetal"),
                18: ("Ar", "Argon", "noble gas"),

                19: ("K", "Potassium", "alkali metal"),
                20: ("Ca", "Calcium", "alkaline earth metal"),
                21: ("Sc", "Scandium", "transition metal"),
                22: ("Ti", "Titanium", "transition metal"),
                23: ("V", "Vanadium", "transition metal"),
                24: ("Cr", "Chromium", "transition metal"),
                25: ("Mn", "Manganese", "transition metal"),
                26: ("Fe", "Iron", "transition metal"),
                27: ("Co", "Cobalt", "transition metal"),
                28: ("Ni", "Nickel", "transition metal"),
                29: ("Cu", "Copper", "transition metal"),
                30: ("Zn", "Zinc", "transition metal"),
                31: ("Ga", "Gallium", "post-transition metal"),
                32: ("Ge", "Germanium", "metalloid"),
                33: ("As", "Arsenic", "metalloid"),
                34: ("Se", "Selenium", "nonmetal"),
                35: ("Br", "Bromine", "nonmetal"),
                36: ("Kr", "Krypton", "noble gas"),

                37: ("Rb", "Rubidium", "alkali metal"),
                38: ("Sr", "Strontium", "alkaline earth metal"),
                39: ("Y", "Yttrium", "transition metal"),
                40: ("Zr", "Zirconium", "transition metal"),
                41: ("Nb", "Niobium", "transition metal"),
                42: ("Mo", "Molybdenum", "transition metal"),
                43: ("Tc", "Technetium", "transition metal"),
                44: ("Ru", "Ruthenium", "transition metal"),
                45: ("Rh", "Rhodium", "transition metal"),
                46: ("Pd", "Palladium", "transition metal"),
                47: ("Ag", "Silver", "transition metal"),
                48: ("Cd", "Cadmium", "transition metal"),
                49: ("In", "Indium", "post-transition metal"),
                50: ("Sn", "Tin", "post-transition metal"),
                51: ("Sb", "Antimony", "metalloid"),
                52: ("Te", "Tellurium", "metalloid"),
                53: ("I", "Iodine", "nonmetal"),
                54: ("Xe", "Xenon", "noble gas"),

                55: ("Cs", "Caesium", "alkali metal"),
                56: ("Ba", "Barium", "alkaline earth metal"),
                # lanthanide
                57: ("La", "Lanthanum", "lanthanide"),
                58: ("Ce", "Cerium", "lanthanide"),
                59: ("Pr", "Praseodymium", "lanthanide"),
                60: ("Nd", "Neodymium", "lanthanide"),
                61: ("Pm", "Promethium", "lanthanide"),
                62: ("Sm", "Samarium", "lanthanide"),
                63: ("Eu", "Europium", "lanthanide"),
                64: ("Gd", "Gadolinium", "lanthanide"),
                65: ("Tb", "Terbium", "lanthanide"),
                66: ("Dy", "Dysprosium", "lanthanide"),
                67: ("Ho", "Holmium", "lanthanide"),
                68: ("Er", "Erbium", "lanthanide"),
                69: ("Tm", "Thulium", "lanthanide"),
                70: ("Yb", "Ytterbium", "lanthanide"),
                71: ("Lu", "Lutetium", "lanthanide"),

                72: ("Hf", "Hafnium", "transition metal"),
                73: ("Ta", "Tantalum", "transition metal"),
                74: ("W", "Tungsten", "transition metal"),
                75: ("Re", "Rhenium", "transition metal"),
                76: ("Os", "Osmium", "transition metal"),
                77: ("Ir", "Iridium", "transition metal"),
                78: ("Pt", "Platinum", "transition metal"),
                79: ("Au", "Gold", "transition metal"),
                80: ("Hg", "Mercury", "transition metal"),
                81: ("Tl", "Thallium", "post-transition metal"),
                82: ("Pb", "Lead", "post-transition metal"),
                83: ("Bi", "Bismuth", "post-transition metal"),
                84: ("Po", "Polonium", "post-transition metal"),
                85: ("At", "Astatine", "metalloid"),
                86: ("Rn", "Radon", "noble gas"),

                87: ("Fr", "Francium", "alkali metal"),
                88: ("Ra", "Radium", "alkaline earth metal"),
                # actinide
                89: ("Ac", "Actinium", "actinide"),
                90: ("Th", "Thorium", "actinide"),
                91: ("Pa", "Protactinium", "actinide"),
                92: ("U", "Uranium", "actinide"),
                93: ("Np", "Neptunium", "actinide"),
                94: ("Pu", "Plutonium", "actinide"),
                95: ("Am", "Americium", "actinide"),
                96: ("Cm", "Curium", "actinide"),
                97: ("Bk", "Berkelium", "actinide"),
                98: ("Cf", "Californium", "actinide"),
                99: ("Es", "Einsteinium", "actinide"),
                100: ("Fm", "Fermium", "actinide"),
                101: ("Md", "Mendelevium", "actinide"),
                102: ("No", "Nobelium", "actinide"),
                103: ("Lr", "Lawrencium", "actinide"),

                104: ("Rf", "Rutherfordium", "transition metal"),
                105: ("Db", "Dubnium", "transition metal"),
                106: ("Sg", "Seaborgium", "transition metal"),
                107: ("Bh", "Bohrium", "transition metal"),
                108: ("Hs", "Hassium", "transition metal"),
            },
            "table_layout": [
                # line 1: H, e⁻, VA, He
                [(1, 1)] + [(None, j) for j in range(2, 7)] + [(-1, 7), (0, 8)] +
                [(None, j) for j in range(9, 18)] + [(2, 18)],
                # line 2
                [(3, 1), (4, 2)] + [(None, j) for j in range(3, 13)] +
                [(5, 13), (6, 14), (7, 15), (8, 16), (9, 17), (10, 18)],
                # line 3
                [(11, 1), (12, 2)] + [(None, j) for j in range(3, 13)] +
                [(j, j) for j in range(13, 19)],
                # line 4
                [(j, j-19+1) for j in range(19, 37)],
                # line 5
                [(j, j-37+1) for j in range(37, 55)],
                # line 6
                [(55, 1), (56, 2), (None, 3)] +
                [(j, j-72+4) for j in range(72, 87)],
                # line 7
                [(87, 1), (88, 2), (None, 3)] +
                [(j, j-104+4) for j in range(104, 109)],
                # line 8, lanthanide
                [(None, j) for j in range(1, 4)] +
                [(j, j-57+4) for j in range(57, 72)],
                # line 9, actinide
                [(None, j) for j in range(1, 4)] +
                [(j, j-89+4) for j in range(89, 104)],
            ]
        }

        # color mapping
        self.color_map = {
            "alkali metal": "#ff9999",
            "alkaline earth metal": "#ffcc99",
            "lanthanide": "#ffccff",
            "actinide": "#ffb3b3",
            "transition metal": "#c2c2d6",
            "post-transition metal": "#ffb366",
            "metalloid": "#ffff99",
            "nonmetal": "#99ff99",
            "noble gas": "#c2f0c2",
            "special": "#FFD700",  # golden for special symbol
            "unknown": "#e0e0e0"
        }

        # init widgets
        self.create_widgets()

    def create_widgets(self):
        """Create all GUI widgets."""
        # 1. TDB file selection
        tk.Label(self.root, text="TDB Database File:").grid(row=0, column=0, sticky="w")
        self.tdb_path = tk.Entry(self.root, width=50)
        self.tdb_path.grid(row=0, column=1)
        tk.Button(self.root, text="Browse", command=self.select_tdb_file).grid(row=0, column=2)

        # 2. Periodic table
        tk.Label(self.root, text="Select Elements and Fractions:").grid(row=1, column=0, sticky="w")
        tk.Label(self.root, text="(Click element buttons below to select, then enter mole fractions in text boxes)").grid(
            row=1, column=1, sticky="w")
        self.create_periodic_table()

        # 3. Phase selection
        tk.Label(self.root, text="Select Phases:").grid(row=5, column=0, sticky="w")
        self.phases_frame = tk.Frame(self.root, relief=tk.SUNKEN, borderwidth=1)
        self.phases_frame.grid(row=6, column=0, columnspan=3, sticky="w")

        # 4. Temperature range
        tk.Label(self.root, text="Temperature Range (K):").grid(row=7, column=0, sticky="w")

        # Create a frame to contain temperature inputs for better layout
        temp_frame = tk.Frame(self.root, relief=tk.RIDGE, borderwidth=1)
        temp_frame.grid(row=7, column=1, columnspan=6, sticky="w")

        tk.Label(temp_frame, text="Start:").grid(row=0, column=0, sticky="e")
        tk.Label(temp_frame, text="End:").grid(row=0, column=2, sticky="e")
        tk.Label(temp_frame, text="Step:").grid(row=0, column=4, sticky="e")
        self.temp_start = tk.Entry(temp_frame, width=8)
        self.temp_end = tk.Entry(temp_frame, width=8)
        self.temp_step = tk.Entry(temp_frame, width=8)
        self.temp_start.insert(0, "273")
        self.temp_end.insert(0, "1773")
        self.temp_step.insert(0, "10")
        self.temp_start.grid(row=0, column=1, sticky="w", padx=2)
        self.temp_end.grid(row=0, column=3, sticky="w", padx=2)
        self.temp_step.grid(row=0, column=5, sticky="w", padx=2)

        # 5. Pressure input
        tk.Label(self.root, text="Pressure (Pa):").grid(row=8, column=0, sticky="w")
        self.pressure = tk.Entry(self.root, width=15)
        self.pressure.insert(0, "100000")
        self.pressure.grid(row=8, column=1, sticky="w")

        # 6. Output directory
        tk.Label(self.root, text="Output Directory:").grid(row=9, column=0, sticky="w")
        self.outdir_path = tk.Entry(self.root, width=50)
        self.outdir_path.insert(0, os.path.join(os.path.dirname(sys.argv[0]), "output"))
        self.outdir_path.grid(row=9, column=1)
        tk.Button(self.root, text="Browse", command=self.select_output_dir).grid(row=9, column=2)

        # 7. Action buttons
        self.calc_button = tk.Button(self.root, text="Calculate", command=self.run_calculation)
        self.calc_button.grid(row=10, column=1, pady=10)
        tk.Button(self.root, text="Exit", command=self.root.quit).grid(row=10, column=2)

        # 8. Log output
        self.log = tk.Text(self.root, height=10, width=70)
        self.log.grid(row=11, column=0, columnspan=3, pady=10)

    def create_periodic_table(self):
        """Create element selection interface based on standard periodic table layout."""
        self.periodic_table_frame = tk.Frame(self.root)
        self.periodic_table_frame.grid(row=2, column=0, columnspan=3)

        # Calculate maximum columns for alignment
        max_cols = max(len(row) for row in self.periodic_table_data["table_layout"])

        for row_idx, row_data in enumerate(self.periodic_table_data["table_layout"]):
            for col_data in row_data:
                if col_data is None:
                    # Blank placeholder
                    label = tk.Label(self.periodic_table_frame, width=5, text="")
                    label.grid(row=row_idx + 1, column=col_data[1] - 1, padx=1, pady=1)
                    continue

                atomic_num, col = col_data
                if atomic_num not in self.periodic_table_data["elements"]:
                    continue

                symbol, name, category = self.periodic_table_data["elements"][atomic_num]
                bg_color = self.color_map.get(category, self.color_map["unknown"])

                btn = tk.Button(
                    self.periodic_table_frame,
                    text=symbol,
                    width=5,
                    bg=bg_color,
                    state=tk.DISABLED,  # Disable all buttons by default
                    command=lambda s=symbol: self.toggle_element(s)
                )
                btn.grid(row=row_idx + 1, column=col - 1, padx=1, pady=1)
                self.element_buttons[symbol] = btn

                # Optional: Add ToolTip to show name
                # tk.ToolTip(btn, text=name)

        # Create frame to display element fractions
        self.elements_input_frame = tk.Frame(self.root)
        self.elements_input_frame.grid(row=3, column=0, columnspan=3, sticky="w")

    def toggle_element(self, element):
        """Toggle element selection status."""
        # Check if TDB file has been loaded
        if not hasattr(self, 'calculator') or self.calculator is None:
            messagebox.showwarning("TDB File Not Loaded", "Please load a TDB file before selecting elements.")
            return

        if element in self.selected_elements:
            # Remove element
            del self.selected_elements[element]
            self.element_buttons[element].config(relief="raised")
        else:
            # Add element, store an initial fraction value (float)
            self.selected_elements[element] = 0.2  # Store float value, not DoubleVar
            self.element_buttons[element].config(relief="sunken")

        self.update_element_inputs()

    def update_element_inputs(self):
        """Update element fraction input fields displayed on the interface."""
        # Clear old input fields
        for widget in self.elements_input_frame.winfo_children():
            widget.destroy()

        # Create new labels and input fields for each selected element
        first_element = True
        first_element_value = 1.0
        elements_list = list(self.selected_elements.items())

        for i, (el, value) in enumerate(elements_list):
            # Element label (special symbols marked with *)
            label_text = f"{el}*" if el in ["VA", "e⁻"] else el
            tk.Label(self.elements_input_frame, text=label_text).grid(row=i, column=0, padx=2, pady=2)

            if first_element and el not in ["VA", "e⁻"]:
                # Create a label instead of an input field to display the auto-calculated value
                # First calculate the sum of other elements
                other_elements_sum = 0
                for other_el, other_value in elements_list[1:]:
                    if other_el not in ["VA", "e⁻"]:
                        other_elements_sum += other_value

                first_element_value = max(0.0, 1.0 - other_elements_sum)  # Ensure non-negative
                # Update the first element's value in the dictionary
                self.selected_elements[el] = first_element_value

                # Display read-only value
                value_label = tk.Label(self.elements_input_frame,
                                       text=f"{first_element_value:.4f}", width=10, bg="#f0f0f0")
                value_label.grid(row=i, column=1, padx=2, pady=2)

                # Add note
                tk.Label(self.elements_input_frame, text="(auto-calculated)",
                         fg="gray").grid(row=i, column=2, padx=2, pady=2)
                first_element = False

            else:
                # For other elements, create editable input fields
                var = tk.DoubleVar(value=value)

                # Capture current el and var to update dictionary when focus is lost
                def make_callback(element_name, double_var):
                    def callback(event):
                        # When input field loses focus,
                        # update DoubleVar value back to self.selected_elements
                        try:
                            new_value = double_var.get()
                            self.selected_elements[element_name] = new_value
                            # Update interface to recalculate first element's value
                            self.update_element_inputs()
                        except tk.TclError:
                            pass
                    return callback

                entry = tk.Entry(self.elements_input_frame, textvariable=var, width=10)
                entry.bind("<FocusOut>", make_callback(el, var))  # Save when losing focus
                entry.grid(row=i, column=1, padx=2, pady=2)

    def update_periodic_table_state(self):
        """Update element availability status based on current TDB file."""
        available_set = set(spec.upper() for spec in self.available_species)

        for symbol, button in self.element_buttons.items():
            if symbol in ["VA", "e⁻"]:
                # Special symbols (VA, e⁻) are always available
                button.config(state=tk.NORMAL, bg="#FFD700")
            else:
                # Check if regular element is in TDB's available species
                if symbol.upper() in available_set:
                    # Get original category to restore color
                    category = next(
                        (cat for num, (sym, _, cat) in
                         self.periodic_table_data["elements"].items()
                         if sym == symbol),
                        "unknown"
                    )
                    bg_color = self.color_map.get(category, self.color_map["unknown"])
                    button.config(state=tk.NORMAL, bg=bg_color)
                else:
                    # Gray indicates unavailable
                    button.config(state=tk.DISABLED, bg="#cccccc")

    def select_tdb_file(self):
        """Select TDB file and initialize calculator."""
        path = filedialog.askopenfilename(filetypes=[("TDB File", "*.tdb")])
        if path:
            self.tdb_path.delete(0, tk.END)
            self.tdb_path.insert(0, path)
            try:
                self.calculator = create_calculator(path)
                self.available_species = [s.upper() for s in self.calculator.get_species()]
                self.update_periodic_table_state()  # Update element availability status
                # Enable buttons for available elements
                for symbol, button in self.element_buttons.items():
                    if button["state"] == tk.NORMAL:
                        button.config(state=tk.NORMAL)
                self.update_phases()
                self.log.insert(tk.END, f"TDB file loaded successfully: {path}\n")
            except Exception as e:
                messagebox.showerror("Load Failed", f"Failed to load TDB file: {str(e)}")
                self.log.insert(tk.END, f"Failed to load TDB file: {str(e)}\n")

    def update_phases(self):
        """Update selectable phase list based on current TDB file."""
        if hasattr(self, "phases_frame"):
            for widget in self.phases_frame.winfo_children():
                widget.destroy()
        self.phase_vars = {}
        if hasattr(self, "calculator") and self.calculator:
            phases = self.calculator.get_phases()
            for i, phase in enumerate(phases):
                var = tk.BooleanVar(value=True)
                cb = tk.Checkbutton(self.phases_frame, text=phase, variable=var)
                cb.grid(row=i // 5, column=i % 5, sticky="w", padx=2, pady=2)
                self.phase_vars[phase] = var

    def select_output_dir(self):
        """Select output directory."""
        path = filedialog.askdirectory()
        if path:
            self.outdir_path.delete(0, tk.END)
            self.outdir_path.insert(0, path)

    def run_calculation(self):
        """Execute CALPHAD equilibrium calculation."""
        try:
            # Parameter collection
            tdb_file = self.tdb_path.get()
            if not os.path.exists(tdb_file):
                raise FileNotFoundError("TDB file does not exist")

            # Phase selection
            selected_phases = [p for p, var in self.phase_vars.items() if var.get()]
            if not selected_phases:
                raise ValueError("No phases selected. Please check at least one phase")

            # Check if elements have been selected
            if not self.selected_elements:
                raise ValueError(
                    "No elements selected. Please select elements")

            # Elements and fractions
            element_fracs = {}
            total = 0

            # Calculate sum of elements other than the first one
            elements_list = list(self.selected_elements.items())
            first_element_key = None
            other_elements_sum = 0

            for i, (el, value) in enumerate(elements_list):
                if el in ["VA", "e⁻"]:
                    # Handle special symbols normally
                    element_fracs[el] = value
                else:
                    if first_element_key is None:
                        # Record the first regular element
                        first_element_key = el
                    else:
                        # Accumulate other elements
                        if not 0.01 <= value <= 0.99:
                            raise ValueError(f"Element {el}'s fraction must be between 0.01-0.99")
                        other_elements_sum += value

            # Validate sum of elements other than the first one
            if first_element_key is not None and not (0.0 <= other_elements_sum <= 0.99):
                raise ValueError(
                    f"Sum of other elements ({other_elements_sum:.2f}) must be between 0.00-0.99 to ensure positive fraction for first element")

            # Set composition of first element (1 - sum of other element compositions)
            if first_element_key is not None:
                first_element_value = 1.0 - other_elements_sum
                element_fracs[first_element_key] = first_element_value
                # Also update value in self.selected_elements for correct display
                self.selected_elements[first_element_key] = first_element_value

            # Add other elements to element_fracs
            for el, value in elements_list:
                if el != first_element_key and el not in ["VA", "e⁻"]:
                    element_fracs[el] = value

            # Temperature parameters
            temp_start = int(self.temp_start.get())
            temp_end = int(self.temp_end.get())
            temp_step = int(self.temp_step.get())

            # Pressure
            pressure = float(self.pressure.get())

            # Output directory
            outdir = self.outdir_path.get()
            if not os.path.exists(outdir):
                os.makedirs(outdir)
            self.log.insert(tk.END, f"Output directory set to: {outdir}\n")

            # Execute calculation
            calculator = create_calculator(tdb_file)
            result = calculator.calculate_equilibrium(
                comps_fracts=element_fracs,
                phases=selected_phases,
                temps=(temp_start, temp_end, temp_step),
                pressure=pressure
            )

            # Plotting
            df_GM = calculator.plot_GM_T(outdir, result)
            df_NP = calculator.plot_NP_T(outdir, result)
            df_X = calculator.plot_X_T(outdir, result)
            df_Y = calculator.plot_Y_T(outdir, result)

            # Write dataframes to Excel file
            with pd.ExcelWriter(os.path.join(outdir, "results.xlsx")) as writer:
                if df_GM is not None and not df_GM.empty:
                    df_GM.to_excel(writer, sheet_name="GM")
                if df_NP is not None and not df_NP.empty:
                    df_NP.to_excel(writer, sheet_name="NP")
                if df_X is not None and not df_X.empty:
                    df_X.to_excel(writer, sheet_name="X")
                if df_Y is not None and not df_Y.empty:
                    df_Y.to_excel(writer, sheet_name="Y")

            self.log.insert(tk.END, f"Calculation completed. Results saved to: {outdir}\n")
            self.log.see(tk.END)  # Auto-scroll to bottom

        except Exception as e:
            messagebox.showerror("Calculation Error", str(e))
            self.log.insert(tk.END, traceback.format_exc() + '\n')
            self.log.insert(tk.END, f"Error: {str(e)}\n")
            self.log.see(tk.END)  # Auto-scroll to bottom


if __name__ == "__main__":
    root = tk.Tk()
    app = CalphadGUI(root)
    root.mainloop()


# nuitka src/gui/guipc.py --standalone --onefile --output-dir=dist `
# --windows-icon-from-ico="src/gui/guipc_icon.ico" `
# --enable-plugin=upx --upx-binary="D:\Programs\upx-5.0.2-win64\upx.exe" `
# --enable-plugin=tk-inter `
# --follow-imports `
# --windows-disable-console
